##独立的KDJ计算器，直接使用xtdata获取当日历史分钟数据
from xtquant import xtdata
from typing import Dict, Optional, Union, Tuple
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import time


class IndependentKDJCalculator:
    """
    完全独立的KDJ计算器

    特性：
    1. 不依赖订阅，直接使用xtdata获取历史数据
    2. 输入股票代码即可获取最新KDJ（上一分钟结束时的KDJ）
    3. 自动下载当日分钟数据并计算
    4. 支持缓存，避免重复下载
    5. 支持预热数据，参考MACD的前一天数据预热方式
    6. KDJ参数：K周期=9，D周期=3，J=3*K-2*D

    使用方法：
        kdj_calc = IndependentKDJCalculator(k_period=9, d_period=3, xtdata=xtdata)
        kdj = kdj_calc.get_latest_kdj("513050.SH")  # 直接获取最新KDJ
    """

    def __init__(self, k_period: int = 9, d_period: int = 3, xtdata=None, use_warmup: bool = True):
        if k_period <= 0 or d_period <= 0:
            raise ValueError("KDJ periods must be positive")
        self.k_period: int = k_period
        self.d_period: int = d_period
        self.xtdata = xtdata
        self.use_warmup = use_warmup

        # 每个股票的状态
        self._last_k: Dict[str, float] = {}
        self._last_d: Dict[str, float] = {}
        self._last_j: Dict[str, float] = {}
        self._last_rsv: Dict[str, float] = {}
        self._price_history: Dict[str, list] = {}  # 存储价格历史用于计算RSV

        # 缓存机制：避免重复下载同一天的数据
        self._cache_date: Dict[str, str] = {}  # 股票代码 -> 缓存日期
        self._cache_ttl: int = 60  # 缓存有效期（秒）- 1分钟更新一次确保获取最新完整分钟数据
        self._cache_time: Dict[str, float] = {}  # 股票代码 -> 缓存时间戳
        self._cache_minute: Dict[str, str] = {}  # 股票代码 -> 缓存的分钟时间戳

        # 交易时间段常量
        self.MORNING_START = datetime.strptime("09:30", "%H:%M").time()
        self.MORNING_END = datetime.strptime("11:30", "%H:%M").time()
        self.AFTERNOON_START = datetime.strptime("13:00", "%H:%M").time()
        self.AFTERNOON_END = datetime.strptime("15:00", "%H:%M").time()

    def reset(self, stock_code: Optional[str] = None) -> None:
        """重置指定股票或全部股票的KDJ状态"""
        if stock_code is None:
            self._last_k.clear()
            self._last_d.clear()
            self._last_j.clear()
            self._last_rsv.clear()
            self._price_history.clear()
            self._cache_date.clear()
            self._cache_time.clear()
            self._cache_minute.clear()
            return

        self._last_k.pop(stock_code, None)
        self._last_d.pop(stock_code, None)
        self._last_j.pop(stock_code, None)
        self._last_rsv.pop(stock_code, None)
        self._price_history.pop(stock_code, None)
        self._cache_date.pop(stock_code, None)
        self._cache_time.pop(stock_code, None)
        self._cache_minute.pop(stock_code, None)

    def _is_cache_valid(self, stock_code: str, today: str) -> bool:
        """检查缓存是否有效"""
        if stock_code not in self._cache_date or stock_code not in self._cache_time:
            return False

        # 检查日期是否匹配
        if self._cache_date[stock_code] != today:
            return False

        # 检查当前分钟是否与缓存分钟相同
        current_minute = datetime.now().strftime("%Y%m%d%H%M")
        cached_minute = self._cache_minute.get(stock_code, "")
        if cached_minute != current_minute:
            return False

        # 检查时间是否在有效期内
        current_time = time.time()
        cache_time = self._cache_time[stock_code]
        return (current_time - cache_time) < self._cache_ttl

    def _download_minute_data(self, stock_code: str, date: str) -> bool:
        """下载指定股票的当日分钟数据"""
        if not self.xtdata:
            print("⚠️ xtdata未初始化，无法下载数据")
            return False

        try:
            self.xtdata.download_history_data(
                stock_code=stock_code,
                period="1m",
                start_time=date,
                end_time=date
            )
            time.sleep(0.1)  # 避免请求过快
            return True
        except Exception as e:
            print(f"⚠️ {stock_code} 分钟数据下载失败: {e}")
            return False

    def _get_minute_data(self, stock_code: str, date: str) -> Optional[pd.DataFrame]:
        """获取指定股票的当日分钟数据"""
        if not self.xtdata:
            return None

        try:
            data = self.xtdata.get_local_data(
                stock_list=[stock_code],
                period='1m',
                start_time=date,
                end_time=date
            )

            df = data.get(stock_code) if data else None
            if df is None or df.empty:
                return None

            # 识别时间列和价格列
            time_col = next((c for c in df.columns if 'time' in c.lower()), None)
            price_col = next((c for c in ['close', 'lastPrice', 'price', 'last', 'current'] if c in df.columns), None)

            if not time_col or not price_col:
                return None

            # 时间格式统一处理
            if str(df[time_col].dtype) in ('int64', 'float64'):
                df['__time__'] = pd.to_datetime(df[time_col], unit='ms') + pd.Timedelta(hours=8)
            else:
                try:
                    df['__time__'] = pd.to_datetime(df[time_col], format='%Y%m%d%H%M%S', errors='coerce')
                except Exception:
                    df['__time__'] = pd.to_datetime(df[time_col], errors='coerce')

            df = df[df['__time__'].notna()].sort_values('__time__')

            # 添加列名到DataFrame以便后续使用
            df._price_col = price_col

            return df if not df.empty else None

        except Exception as e:
            print(f"⚠️ {stock_code} 分钟数据获取失败: {e}")
            return None

    def _filter_trading_session_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """筛选交易时段内的数据，并排除当前未完成的分钟"""
        try:
            # 先筛选交易时段内的数据
            mask = [((self.MORNING_START <= ts.time() <= self.MORNING_END) or
                     (self.AFTERNOON_START <= ts.time() <= self.AFTERNOON_END)) for ts in df['__time__']]
            if any(mask):
                filtered_df = df[mask]

                # 排除当前未完成的分钟：只保留完整的分钟数据
                current_time = datetime.now()
                current_minute = current_time.replace(second=0, microsecond=0)

                # 只保留时间小于当前分钟的数据（即已完成的分钟）
                complete_mask = filtered_df['__time__'] < current_minute
                if complete_mask.any():
                    return filtered_df[complete_mask]
                else:
                    return filtered_df  # 如果没有完整分钟数据，返回原数据
            else:
                return df
        except Exception:
            pass
        return df

    def get_previous_day_ticks_for_warmup(self, stock_code: str) -> Optional[pd.DataFrame]:
        """
        获取前一天的tick数据用于预热，只取下午1点到3点的数据
        参考MACD的预热数据获取方式

        Args:
            stock_code: 股票代码

        Returns:
            过滤后的前一天下午tick数据DataFrame
        """
        try:
            # 获取前一个交易日的日期（如果是周一则获取上周五）
            today = datetime.now()
            if today.weekday() == 0:  # 周一
                # 获取上周五的日期
                yesterday = (today - timedelta(days=3)).strftime("%Y%m%d")
            else:
                # 获取前一天的日期
                yesterday = (today - timedelta(days=2)).strftime("%Y%m%d")

            # 下载前一天的tick数据
            self.xtdata.download_history_data(stock_code=stock_code, period="tick",
                                              start_time=yesterday, end_time=yesterday)
            time.sleep(2)

            # 获取本地数据
            data = self.xtdata.get_local_data(stock_list=[stock_code], period='tick',
                                              start_time=yesterday, end_time=yesterday)

            if not data or stock_code not in data:
                return None

            df = data[stock_code]
            if df.empty:
                return None

            # 确定时间列和价格列
            time_col = next((col for col in df.columns if 'time' in col.lower()), None)
            price_col = next((col for col in ['price', 'lastPrice', 'close', 'last', 'current']
                              if col in df.columns), None)

            if not time_col or not price_col:
                return None

            # 转换时间戳
            try:
                if df[time_col].dtype in ['int64', 'float64']:
                    df['time'] = pd.to_datetime(df[time_col], unit='ms') + pd.Timedelta(hours=8)
                else:
                    df['time'] = pd.to_datetime(df[time_col], format='%Y%m%d%H%M%S', errors='coerce')

                df = df[df['time'].notna()].sort_values('time')
                if df.empty:
                    return None
            except Exception:
                return None

            # 过滤无效价格数据
            df_filtered = df[df[price_col] > 0].copy()
            if df_filtered.empty:
                return None

            # 只保留下午1点到3点的数据作为预热数据
            afternoon_start = pd.Timestamp('13:00:00').time()
            afternoon_end = pd.Timestamp('15:00:00').time()

            df_filtered['time_only'] = df_filtered['time'].dt.time
            df_filtered = df_filtered[
                (df_filtered['time_only'] >= afternoon_start) & (df_filtered['time_only'] <= afternoon_end)
                ]

            # 调试信息：显示预热数据的时间范围
            if not df_filtered.empty:
                print(f"预热数据时间范围: {df_filtered['time'].min().strftime('%H:%M')} 到 {df_filtered['time'].max().strftime('%H:%M')}")
                print(f"预热数据点数量: {len(df_filtered)}")

            if df_filtered.empty:
                return None

            # 只保留必要的列
            result_df = df_filtered[['time', price_col]].copy()
            result_df.rename(columns={price_col: 'close'}, inplace=True)

            return result_df

        except Exception as e:
            print(f"获取前一天tick数据失败: {e}")
            return None

    def resample_ticks_to_minutes(self, df: pd.DataFrame, price_col: str = 'close') -> pd.DataFrame:
        """将tick数据重新采样为分钟数据"""
        if 'time' not in df.columns:
            return pd.DataFrame()

        df_resampled = df.set_index('time')
        df_minutes = df_resampled[price_col].resample('1min').agg({
            'open': 'first',
            'high': 'max',
            'low': 'min',
            'close': 'last',
            'volume': 'sum' if 'volume' in df_resampled.columns else lambda x: len(x)
        }).dropna()
        return df_minutes.reset_index()

    def _calculate_rsv(self, high_prices: list, low_prices: list, close_prices: list, period: int) -> float:
        """计算RSV值"""
        if len(close_prices) < period:
            return 50.0  # 默认值

        # 取最近period个数据点
        recent_highs = high_prices[-period:]
        recent_lows = low_prices[-period:]
        current_close = close_prices[-1]

        highest_high = max(recent_highs)
        lowest_low = min(recent_lows)

        if highest_high == lowest_low:
            return 50.0  # 避免除零

        rsv = ((current_close - lowest_low) / (highest_high - lowest_low)) * 100
        return rsv

    def _calculate_kdj_from_data(self, stock_code: str, df: pd.DataFrame, warmup_data: Optional[pd.DataFrame] = None) -> Optional[Dict[str, float]]:
        """从分钟数据计算KDJ"""
        try:
            # 安全地获取价格列名
            price_col = getattr(df, '_price_col', None)
            if not price_col:
                # 如果没有_price_col属性，重新识别价格列
                price_col = next((c for c in ['close', 'lastPrice', 'price', 'last', 'current'] if c in df.columns),
                                 None)
                if not price_col:
                    print(f"⚠️ {stock_code} 无法识别价格列")
                    return None

            # 准备价格数据
            prices = df[price_col].dropna().tolist()
            if len(prices) < 2:
                return None

            # 处理预热数据
            warmup_prices = []
            if warmup_data is not None and not warmup_data.empty:
                warmup_minutes = self.resample_ticks_to_minutes(warmup_data, 'close')
                if not warmup_minutes.empty:
                    warmup_prices = warmup_minutes['close'].tolist()

            # 合并预热数据和当日数据
            all_prices = warmup_prices + prices
            if len(all_prices) < self.k_period:
                return None

            # 重置该股票的状态
            self._last_k[stock_code] = 50.0  # 初始K值
            self._last_d[stock_code] = 50.0  # 初始D值
            self._last_j[stock_code] = 50.0  # 初始J值
            self._price_history[stock_code] = []

            # 依次处理每个价格
            for i, price in enumerate(all_prices):
                self._price_history[stock_code].append(price)
                
                # 计算RSV
                if len(self._price_history[stock_code]) >= self.k_period:
                    # 使用当前价格作为high、low、close（简化处理）
                    high_prices = self._price_history[stock_code]
                    low_prices = self._price_history[stock_code]
                    close_prices = self._price_history[stock_code]
                    
                    rsv = self._calculate_rsv(high_prices, low_prices, close_prices, self.k_period)
                    self._last_rsv[stock_code] = rsv
                    
                    # 更新K、D、J值
                    self._update_kdj_state(stock_code, rsv)

            return {
                'k': self._last_k.get(stock_code, 50.0),
                'd': self._last_d.get(stock_code, 50.0),
                'j': self._last_j.get(stock_code, 50.0),
                'rsv': self._last_rsv.get(stock_code, 50.0)
            }

        except Exception as e:
            print(f"⚠️ {stock_code} KDJ计算异常: {e}")
            return None

    def _update_kdj_state(self, stock_code: str, rsv: float) -> None:
        """更新KDJ内部状态"""
        rsv = float(rsv)

        # K值计算：K = (2/3) * 前K + (1/3) * 当前RSV
        if stock_code not in self._last_k:
            self._last_k[stock_code] = 50.0
        else:
            self._last_k[stock_code] = (2.0/3.0) * self._last_k[stock_code] + (1.0/3.0) * rsv

        # D值计算：D = (2/3) * 前D + (1/3) * 当前K
        if stock_code not in self._last_d:
            self._last_d[stock_code] = 50.0
        else:
            self._last_d[stock_code] = (2.0/3.0) * self._last_d[stock_code] + (1.0/3.0) * self._last_k[stock_code]

        # J值计算：J = 3*K - 2*D
        self._last_j[stock_code] = 3.0 * self._last_k[stock_code] - 2.0 * self._last_d[stock_code]

    def _filter_data_before_time(self, df: pd.DataFrame, end_time: datetime) -> pd.DataFrame:
        """
        过滤掉指定时间之后的数据（回测模式使用）

        Args:
            df: 原始数据DataFrame
            end_time: 结束时间

        Returns:
            过滤后的DataFrame
        """
        try:
            if df.empty:
                return df

            # 确保时间列存在
            if '__time__' not in df.columns:
                return df

            # 过滤掉结束时间之后的数据
            filtered_df = df[df['__time__'] <= end_time]
            return filtered_df

        except Exception as e:
            print(f"⚠️ 数据过滤异常: {e}")
            return df

    def get_latest_kdj(self, stock_code: str, end_time: Optional[datetime] = None) -> Optional[Dict[str, float]]:
        """
        获取指定股票的最新KDJ值（上一分钟结束时的KDJ）

        Args:
            stock_code: 股票代码
            end_time: 结束时间（回测模式使用，实时模式为None）

        Returns:
            KDJ值字典，包含k、d、j、rsv，如果无法计算则返回None
        """
        try:
            # 确定当前日期
            if end_time is not None:
                # 回测模式：使用指定的结束时间
                today = end_time.strftime("%Y%m%d")
                current_time = end_time
            else:
                # 实时模式：使用当前时间
                today = datetime.now().strftime("%Y%m%d")
                current_time = datetime.now()

            # 检查缓存是否有效
            if self._is_cache_valid(stock_code, today):
                return {
                    'k': self._last_k.get(stock_code, 50.0),
                    'd': self._last_d.get(stock_code, 50.0),
                    'j': self._last_j.get(stock_code, 50.0),
                    'rsv': self._last_rsv.get(stock_code, 50.0)
                }

            # 缓存无效，需要重新下载和计算
            print(f"🔄 {stock_code} 开始下载当日分钟数据...")

            # 下载分钟数据
            if not self._download_minute_data(stock_code, today):
                return None

            # 获取分钟数据
            df = self._get_minute_data(stock_code, today)
            if df is None:
                return None

            # 筛选交易时段内的数据
            df = self._filter_trading_session_data(df)
            if df.empty:
                return None

            # 在回测模式下，过滤掉结束时间之后的数据
            if end_time is not None:
                df = self._filter_data_before_time(df, end_time)
                if df.empty:
                    return None

            # 获取预热数据
            warmup_data = None
            if self.use_warmup:
                try:
                    warmup_data = self.get_previous_day_ticks_for_warmup(stock_code)
                except Exception as e:
                    print(f"预热数据获取失败，使用无预热模式: {e}")
                    warmup_data = None

            # 计算KDJ
            kdj_result = self._calculate_kdj_from_data(stock_code, df, warmup_data)

            # 更新缓存
            if kdj_result is not None:
                self._cache_date[stock_code] = today
                self._cache_time[stock_code] = time.time()
                self._cache_minute[stock_code] = current_time.strftime("%Y%m%d%H%M")
                print(f"✅ {stock_code} KDJ计算完成: K={kdj_result['k']:.2f}, D={kdj_result['d']:.2f}, J={kdj_result['j']:.2f}")
            else:
                print(f"⚠️ {stock_code} KDJ计算失败")

            return kdj_result

        except Exception as e:
            print(f"❌ {stock_code} 获取KDJ异常: {e}")
            return None

    def get_kdj_with_preview(self, stock_code: str, current_price: Optional[float] = None) -> Dict[str, Optional[Dict[str, float]]]:
        """
        获取KDJ值，包括确认值和预览值

        Args:
            stock_code: 股票代码
            current_price: 当前价格（用于预览KDJ）

        Returns:
            包含确认KDJ和预览KDJ的字典
        """
        # 获取确认KDJ
        confirmed_kdj = self.get_latest_kdj(stock_code)

        # 计算预览KDJ
        preview_kdj = None
        if current_price is not None and confirmed_kdj is not None:
            preview_kdj = self._calculate_preview_kdj(stock_code, current_price)

        return {
            'confirmed': confirmed_kdj,
            'preview': preview_kdj
        }

    def _calculate_preview_kdj(self, stock_code: str, current_price: float) -> Optional[Dict[str, float]]:
        """计算预览KDJ（基于当前价格）"""
        try:
            if stock_code not in self._price_history or len(self._price_history[stock_code]) < self.k_period:
                return None

            # 添加当前价格到历史
            temp_history = self._price_history[stock_code] + [current_price]
            
            # 计算预览RSV
            high_prices = temp_history
            low_prices = temp_history
            close_prices = temp_history
            
            preview_rsv = self._calculate_rsv(high_prices, low_prices, close_prices, self.k_period)
            
            # 计算预览K、D、J
            current_k = self._last_k.get(stock_code, 50.0)
            current_d = self._last_d.get(stock_code, 50.0)
            
            preview_k = (2.0/3.0) * current_k + (1.0/3.0) * preview_rsv
            preview_d = (2.0/3.0) * current_d + (1.0/3.0) * preview_k
            preview_j = 3.0 * preview_k - 2.0 * preview_d

            return {
                'k': preview_k,
                'd': preview_d,
                'j': preview_j,
                'rsv': preview_rsv
            }

        except Exception:
            return None

    def force_refresh(self, stock_code: str) -> Optional[Dict[str, float]]:
        """强制刷新指定股票的KDJ数据"""
        self.reset(stock_code)
        return self.get_latest_kdj(stock_code)

    def get_all_stock_kdj(self, stock_codes: list) -> Dict[str, Optional[Dict[str, float]]]:
        """批量获取多个股票的KDJ值"""
        results = {}
        for code in stock_codes:
            results[code] = self.get_latest_kdj(code)
        return results


# 为了保持向后兼容，保留原有的类名
StandardKDJCalculator = IndependentKDJCalculator

# 更新导出列表
__all__ = ["IndependentKDJCalculator", "StandardKDJCalculator"]

# 使用示例
if __name__ == "__main__":
    # 定义股票代码
    stock_code = "513750.SH"  # 可以修改为您需要的股票代码
    
    # 创建KDJ计算器（K周期=9，D周期=3）
    kdj_calc = IndependentKDJCalculator(k_period=9, d_period=3, xtdata=xtdata, use_warmup=True)
    
    # 获取最新KDJ
    target_time = datetime(2025, 9, 3, 11, 10, 1)  # 14:30:00
    kdj_value = kdj_calc.get_latest_kdj(stock_code, target_time)
    
    if kdj_value is not None:
        print(f"股票 {stock_code} 的最新KDJ:")
        print(f"  K值: {kdj_value['k']:.2f}")
        print(f"  D值: {kdj_value['d']:.2f}")
        print(f"  J值: {kdj_value['j']:.2f}")
        print(f"  RSV: {kdj_value['rsv']:.2f}")
    else:
        print(f"无法获取 {stock_code} 的KDJ值")
